import argparse
import os
import sys

sys.path.append(os.getcwd())

from pytorch_lightning import Trainer
from pytorch_lightning import loggers as pl_loggers
from pytorch_lightning.callbacks import ModelCheckpoint

from utils.config import Config
from utils.sys import YoloSystem


def parse_args():
    argparser = argparse.ArgumentParser()
    argparser.add_argument("--devices", type=int, nargs='+', required=False)
    args = argparser.parse_args()

    config = Config()
    config.load("configs/yolo.json")

    if args.devices is not None:
        config.TRAIN.DEVICES = args.devices

    return config


if __name__ == "__main__":
    config = parse_args()

    tb_logger = pl_loggers.TensorBoardLogger(".", name="logs")
    os.mkdir(tb_logger.log_dir)
    config.dump(os.path.join(tb_logger.log_dir, "config.json"))

    # 定义checkpoint回调
    checkpoint_callback = ModelCheckpoint(
        monitor="epoch",
        dirpath=tb_logger.log_dir,
        filename="model-checkpoint",
        save_last=True,
        save_weights_only=False,
        save_top_k=1,
    )
    final_chaeckpoint_callback = ModelCheckpoint(
        monitor="val/loss",
        dirpath=tb_logger.log_dir,
        filename="best-model",
        save_top_k=1,
        save_weights_only=True,
    )

    # 实例化Trainer并开始训练
    trainer = Trainer(
        max_epochs=config.TRAIN.EPOCHS,
        callbacks=[checkpoint_callback, final_chaeckpoint_callback],
        logger=tb_logger,
        log_every_n_steps=1,
        devices=config.TRAIN.DEVICES,
    )
    with trainer.init_module():
        model = YoloSystem(
            backbone=config.MODEL.BACKBONE,
            lambda_coord=config.LOSS.LAMBDA_COORD,
            lambda_noobj=config.LOSS.LAMBDA_NOOBJ,
            grid_size=config.MODEL.GRID_SIZE,
            bounding_box_num=config.MODEL.BOUNDING_BOX_NUM,
            data_path=config.DATASET.ROOT_PATH,
            data_series=config.DATASET.SERIES,
            data_version=config.DATASET.VERSION,
            auto_download=config.DATASET.AUTO_DOWNLOAD,
            batch_size=config.TRAIN.BATCH_SIZE,
            workers_num=config.TRAIN.WORKERS_NUM,
        )
    trainer.fit(model)
